import cv2
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit

# 超参数配置
BATCH_SIZE = 1
INPUT_HW = (640, 640)  # 与导出时的尺寸一致
CLASS_NAMES = ["person", "bicycle", "car", ...]  # 替换为你的类别标签

# --------------------------------------
# **1. 加载 TensorRT 引擎**
# --------------------------------------
logger = trt.Logger(trt.Logger.WARNING)
runtime = trt.Runtime(logger)

with open("yolov5s.engine", "rb") as f:
    serialized_engine = f.read()
engine = runtime.deserialize_cuda_engine(serialized_engine)
context = engine.create_execution_context()

# 获取输入输出绑定信息
input_binding = engine[0]  # 假设只有一个输入
output_binding = engine[1]  # 根据实际输出数量调整

# 分配 GPU 内存
input_host = cuda.pagelocked_empty(trt.volume(input_binding.shape), dtype=np.float32)
input_device = cuda.mem_alloc(input_host.nbytes)
output_host = cuda.pagelocked_empty(trt.volume(output_binding.shape), dtype=np.float32)
output_device = cuda.mem_alloc(output_host.nbytes)
stream = cuda.Stream()


# --------------------------------------
# **2. 图像预处理 (关键步骤)**
# --------------------------------------
def preprocess(image):
    # Letterbox 缩放 (保持宽高比)
    h, w = image.shape[:2]
    scale = min(INPUT_HW[1] / h, INPUT_HW[0] / w)
    nh, nw = int(h * scale), int(w * scale)
    image_resized = cv2.resize(image, (nw, nh))

    # 创建画布并填充边缘
    canvas = np.full((INPUT_HW[1], INPUT_HW[0], 3), 114, dtype=np.uint8)
    dh, dw = (INPUT_HW[1] - nh) // 2, (INPUT_HW[0] - nw) // 2
    canvas[dh:dh + nh, dw:dw + nw] = image_resized

    # 归一化并转置为 CHW 格式
    blob = canvas.astype(np.float32) / 255.0
    blob = blob.transpose(2, 0, 1)[np.newaxis, ...]  # [1,3,640,640]
    return blob, (scale, dw, dh)


# --------------------------------------
# **3. 执行推理**
# --------------------------------------
def infer(image):
    # 预处理
    blob, (scale, dw, dh) = preprocess(image)
    np.copyto(input_host, blob.ravel())

    # 将数据从 CPU 复制到 GPU
    cuda.memcpy_htod_async(input_device, input_host, stream)

    # 执行推理
    context.execute_async_v2(
        bindings=[int(input_device), int(output_device)],
        stream_handle=stream.handle
    )

    # 将结果从 GPU 复制回 CPU
    cuda.memcpy_dtoh_async(output_host, output_device, stream)
    stream.synchronize()

    # 输出形状处理 (假设输出为 [1,25200,85])
    outputs = output_host.reshape(1, 25200, -1)
    return outputs, (scale, dw, dh)


# --------------------------------------
# **4. 后处理 (NMS 和坐标转换)**
# --------------------------------------
def postprocess(outputs, conf_thres=0.5, iou_thres=0.5):
    # 提取检测结果
    detections = outputs[0]  # [25200,85]

    # 过滤低置信度检测框
    mask = detections[..., 4] > conf_thres
    detections = detections[mask]

    # 转换坐标 (xywh to xyxy)
    boxes = detections[..., :4].copy()
    boxes[..., 0] = (boxes[..., 0] - dw) / scale  # x 方向还原
    boxes[..., 1] = (boxes[..., 1] - dh) / scale  # y 方向还原
    boxes[..., 2] = boxes[..., 2] / scale  # w 还原
    boxes[..., 3] = boxes[..., 3] / scale  # h 还原
    boxes = xywh2xyxy(boxes)  # 转换为左上右下坐标

    # 非极大值抑制 (NMS)
    scores = detections[..., 4]
    class_ids = np.argmax(detections[..., 5:], axis=-1)
    indices = cv2.dnn.NMSBoxes(boxes.tolist(), scores.tolist(), conf_thres, iou_thres)

    return [boxes[indices], scores[indices], class_ids[indices]]


def xywh2xyxy(boxes):
    """将 [x_center, y_center, w, h] 转换为 [x1, y1, x2, y2]"""
    return np.column_stack((
        boxes[:, 0] - boxes[:, 2] / 2,
        boxes[:, 1] - boxes[:, 3] / 2,
        boxes[:, 0] + boxes[:, 2] / 2,
        boxes[:, 1] + boxes[:, 3] / 2
    ))


# --------------------------------------
# **5. 使用示例**
# --------------------------------------
if __name__ == "__main__":
    # 读取图像
    image = cv2.imread("test.jpg")

    # 推理
    outputs, meta = infer(image)
    boxes, scores, class_ids = postprocess(outputs, meta)

    # 可视化结果
    for box, score, cls_id in zip(boxes, scores, class_ids):
        x1, y1, x2, y2 = map(int, box)
        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
        label = f"{CLASS_NAMES[cls_id]}: {score:.2f}"
        cv2.putText(image, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

    cv2.imwrite("result.jpg", image)